import sys
import zlib
from zipfile import ZipFile
import os
import xml.etree.ElementTree as ET
import re

XML_HEADER = '''<?xml version="1.0" encoding="UTF-8" ?>
'''

class ParserData:
	def __init__(self, data_root):
		self.data_root = data_root
	
	# Get an absolute path to a folder
	def get_path(self, *paths):
		args = (self.data_root,) + paths
		return os.path.join(*args)
	
	# List file in a source folder
	def list_files(self, *paths):
		source_folder = self.get_path(*paths)
		all_files = [ f for f in os.listdir(source_folder) if os.path.isfile(os.path.join(source_folder, f)) ]
		return all_files

	# Read a file content
	def read_file(self, *paths):
		file_path = self.get_path(*paths)
		content = ''
		with open(file_path, 'rb') as a_file:
			content = a_file.read()
		return content
		
# Split a file into many XML files
def split_file(content):
	start = 0
	start_location = content.find('<DOC>', start)
	end_location = content.find('</DOC>', start_location)
	parts = [] # a 
	while start_location > -1 and end_location > -1:
		part = XML_HEADER + content[start_location:end_location + 6]
		parts.append(part)
		# next part ...
		start = end_location
		start_location = content.find('<DOC>', start)
		end_location = content.find('</DOC>', start_location)
	return parts

def find_entity(content, verbose=True):
	items = []
	allow_entities = [ '&gt;', '&amp;', '&quot;', '&lt;', '&gt;']
	entity_regex = re.compile('&[a-zA-Z0-9]+;')
	matches = entity_regex.findall(content)
	if matches:
		for match in matches:
			if match not in allow_entities and match not in items:
				items.append(match)
				if verbose:
					print(match)
	return items

def fix_sgml(content):
	content = re.sub(r'(<[A-Z] +[A-Z]=)([0-9a-zA-Z]+)(>)',r'\1"\2"\3', content)
	#("&amp;", "&gt;", "&less;") ()
	content = content.replace('&less;', '&lt;')
	
	
	'''
	quot 	" 	U+0022 (34) 	XML 1.0 	double quotation mark
amp 	& 	U+0026 (38) 	XML 1.0 	ampersand
apos 	' 	U+0027 (39) 	XML 1.0 	apostrophe (= apostrophe-quote)
lt 	< 	U+003C (60) 	XML 1.0 	less-than sign
gt
	'''
	
	
	entities = find_entity(content, False)
	for enti in entities:
		content = content.replace(enti, '')
	
	for sgml_entity in ["&M;", "@O;", "&P;", "&rsqb;", "&lsqb;", "&plus;", "&equals;"]:
		content = content.replace(sgml_entity, '')
	return content
# Zip data
def convert_then_zip_data(parser, folder_name, limit=-1):
	text_files = parser.list_files(folder_name)
	with ZipFile(parser.get_path(folder_name + '.zip'), 'w') as package_file:
		i = 0
		for idx in range(len(text_files)):
			text_file = text_files[idx]
			content = parser.read_file(folder_name, text_file)
			content = fix_sgml(content)
			print("Processing: %s (%d of %d)" % (text_file, idx + 1, len(text_files))),
			parts = split_file(content)
			print(" ... Saving %04d entries ..." % len(parts))
			for i in range(len(parts)):
				entry_name = "%s-%04d.xml" % (text_file, i+1)
				package_file.writestr(entry_name, parts[i])
			i += 1
			if limit > 0 and limit <= i:
				break
			#break
	print("DONE!")

# Show information of a zip corpus
def show_info(parser, file_name):
	file_path = parser.get_path(file_name)
	print("Package path: %s" % file_path)
	with ZipFile(file_path, 'r') as package_file:
		names = package_file.namelist()
		print("Entries: %d" % len(names))

def list_content(parser, file_name, verbose=True):
	file_path = parser.get_path(file_name)
	print("Package path: %s" % file_path)
	with ZipFile(file_path, 'r') as package_file:
		names = package_file.namelist()
		if verbose:
			for name in names:
				print(name + "\t"),
	return names
		
# Read an entry in zipped package
def read_content(parser, file_name, entry_name):
	file_path = parser.get_path(file_name)
	with ZipFile(file_path, 'r') as package_file:
		content = package_file.read(entry_name)
	return content

'''
#PROFILES:
AP : title = <HEAD> - text = <TEXT>
CR_EFILES: title = <TEXT>\<TTL> - text = <TEXT>
CR_HFILES:
DOE
FBIS
FT
LATIMES
PATENTS
SJM
WSJ
ZIFF
'''
	
# Read the text content from XML content
def read_text_content(xml_content):
	root = ET.fromstring(xml_content)
	text_element = root.find('TEXT')
	#for child in text_element:
	#	print "[%s]" % child.tail.strip()
	#return text_element.tail
	return ET.tostring(text_element)
	
def print_usage():
	print('''TREC TOOLS
Usage: 
	trectool.py zip folder_name
	trectool.py info zip_file
	trectool.py list zip_file
	trectool.py read zip_file entry_name
''')
	
if __name__ == "__main__":
	parser = ParserData(data_root = 'C:/ca1/raw')
	param_count = len(sys.argv)
	
	# Control
	if param_count <= 1:
		print_usage()
	elif param_count == 3 and sys.argv[1] == 'zip':
		folder_name = sys.argv[2]
		convert_then_zip_data(parser, folder_name)
	elif param_count == 4 and sys.argv[1] == 'zip':
		folder_name = sys.argv[2]
		limit = int(sys.argv[3])
		convert_then_zip_data(parser, folder_name, limit=limit)
	elif sys.argv[1] == 'info':
		if param_count == 3:
			file_name = sys.argv[2]
			show_info(parser, file_name)
		else:
			# display stat of all zipped packages
			files = parser.list_files()
			for a_file in files:
				if a_file.endswith(".zip"):
					show_info(parser, a_file)
	elif sys.argv[1] == 'list':
		if param_count == 3:
			file_name = sys.argv[2]
			list_content(parser, file_name)
	elif sys.argv[1] == 'read' and param_count == 4: #read an entry in zip package
		file_name = sys.argv[2]
		entry_name = sys.argv[3]
		content = read_content(parser, file_name, entry_name)
		print(read_text_content(content))
	else:
		print_usage()
